package com.lizba.redis.limit.funnel;

import redis.clients.jedis.Jedis;

import java.util.HashMap;
import java.util.Map;

/**
 * <p>
 *      redis hash 漏斗限流
 * </p>
 *
 * @Author: Liziba
 * @Date: 2021/9/7 23:46
 */
public class FunnelRateLimiterByHash {

    private Jedis client;

    public FunnelRateLimiterByHash(Jedis client) {
        this.client = client;
    }

    /**
     * 请求是否成功
     *
     * @param userId
     * @param actionKey
     * @param capacity
     * @param leakingRate
     * @param quota
     * @return
     */
    public boolean isActionAllowed(String userId, String actionKey, int capacity, float leakingRate, int quota) {
        String key = this.key(userId, actionKey);
        long nowTs = System.currentTimeMillis();
        Map<String, String> funnelMap = client.hgetAll(key);
        if (funnelMap == null || funnelMap.isEmpty()) {
            return initFunnel(key, nowTs, capacity, quota);
        }
        long intervalTs = nowTs - Long.parseLong(funnelMap.get("leakingTs"));
        int intervalCapacity = (int) (intervalTs * leakingRate);
        // 时间过长, int可能溢出
        if (intervalCapacity < 0) {
            intervalCapacity = 0;
            initFunnel(key, nowTs, capacity, quota);
        }
        // 腾出空间必须 >= 1
        if (intervalCapacity < 1) {
            intervalCapacity = 0;
        }
        int leftCapacity = Integer.parseInt(funnelMap.get("leftCapacity")) + intervalCapacity;
        if (leftCapacity > capacity) {
            leftCapacity = capacity;
        }
        return initFunnel(key, nowTs, leftCapacity, quota);
    }



    /**
     * 存入redis，初始funnel
     *
     * @param key
     * @param nowTs
     * @param capacity
     * @param quota
     * @return
     */
    private boolean initFunnel(String key,long nowTs, int capacity, int quota) {
        Map<String, String> funnelMap = new HashMap<>();
        funnelMap.put("leftCapacity", String.valueOf((capacity > quota) ? (capacity - quota) : 0));
        funnelMap.put("leakingTs", String.valueOf(nowTs));
        client.hset(key, funnelMap);
        return capacity >= quota;
    }

    /**
     * 限流key
     *
     * @param userId
     * @param actionKey
     * @return
     */
    private String key(String userId, String actionKey) {
        return String.format("limit:%s:%s", userId, actionKey);
    }
    
}
